package com.ai.aiagent.advisor;


import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
import reactor.core.publisher.Flux;

import java.util.List;

/**
 * 自定义违禁词advisor
 */
@Slf4j
public class WordsCheckAdvisor  implements CallAdvisor, StreamAdvisor {

    private final int order;
    private static final List<String> bannedWords = List.of("色情", "政治", "暴力");

    public WordsCheckAdvisor(int o) {
        this.order = o;
    }

    public void wordCheck(ChatClientRequest chatClientRequest) {
        // 权限校验逻辑示例：从请求中获取用户信息进行校验
        String userMessage = chatClientRequest.prompt().getUserMessage().getText();
        StringBuilder bannedWordsFound = new StringBuilder();
        
        bannedWords.stream()
                .filter(userMessage::contains)
                .forEach(word -> {
                    if (bannedWordsFound.length() > 0) {
                        bannedWordsFound.append(", ");
                    }
                    bannedWordsFound.append(word);
                });
        
        if (bannedWordsFound.length() > 0) {
            log.error("用户信息包含敏感词: {}", bannedWordsFound.toString());
            throw new RuntimeException("用户信息包含以下敏感词: " + bannedWordsFound.toString());
        }
    }

    /**
     * 权限校验
     * @param chatClientRequest
     * @param callAdvisorChain
     * @return
     */
    @Override
    public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain callAdvisorChain) {
        wordCheck(chatClientRequest);
        return callAdvisorChain.nextCall(chatClientRequest);
    }

    /**
     * 流式敏感词校验
     * @param chatClientRequest
     * @param streamAdvisorChain
     * @return
     */
    @Override
    public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest, StreamAdvisorChain streamAdvisorChain) {
        wordCheck(chatClientRequest);
        return streamAdvisorChain.nextStream(chatClientRequest);
    }


    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }

    /**
     * 权限校验的优先级最高
     * @return
     */
    @Override
    public int getOrder() {
        return this.order;
    }

    public static Builder builder() {
        return new Builder();
    }

    public static final class Builder {
        private int order = 0;

        private Builder() {
        }

        public Builder order(int order) {
            this.order = order;
            return this;
        }

        public WordsCheckAdvisor build() {
            return new WordsCheckAdvisor(order);
        }
    }
}
